modelopt/onnx/utils.py (1)

1218-1261: Consider edge case where first Cast has multiple consumers.

The function checks len(node.outputs[0].outputs) != 1 (line 1231) to ensure the first Cast's output goes to exactly one node. However, this may be overly restrictive. If the first Cast feeds into a duplicate second Cast AND other nodes, you could still remove the duplicate Cast while preserving the connection to other consumers. The current logic skips this optimization opportunity.

This is a minor optimization opportunity and the current implementation is safe.

modelopt/torch/_deploy/utils/torch_onnx.py (1)

62-71: Scope this onnxconverter_common workaround to the conversion call.

Patching the module at import time changes behavior process-wide, and suppress(AttributeError) hides every upstream AttributeError, not just the known list/attr bug. A temporary patch around convert_float_to_float16() is much safer, and it avoids making this module import brittle if the upstream symbol changes.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/_deploy/utils/torch_onnx.py` around lines 62 - 71, The current
import-time monkey-patch of _f16_module.remove_unnecessary_cast_node (using
_original_remove_unnecessary_cast_node and
_patched_remove_unnecessary_cast_node) is global and hides all AttributeError
via suppress(AttributeError); instead, scope the workaround only around the call
to convert_float_to_float16(): before calling convert_float_to_float16() save
the original _f16_module.remove_unnecessary_cast_node, replace it with a minimal
wrapper that only catches the specific list/attribute error, call
convert_float_to_float16(), and finally restore the original function in a
try/finally so the patch is temporary and does not swallow unrelated
AttributeErrors or affect the rest of the process.

modelopt/onnx/utils.py (1)

1507-1511: ⚠️ Potential issue | 🟠 Major

Guard FP32→FP16 cast rewrites by input dtype.

At Line 1507, this rewrites any Cast(to=FLOAT) feeding target ops, including casts from non-FP32 sources. That can silently change behavior by removing intentional upcasts.

💡 Proposed fix

     for node in model.graph.node:
         if node.op_type != "Cast":
             continue
@@
         if not feeds_target:
             continue
 
+        cast_input_type = _get_tensor_type_by_name(model, node.input[0])
+        if cast_input_type != onnx.TensorProto.FLOAT:
+            continue
+
         # Check if Cast is to FP32, and change to FP16
         for attr in node.attribute:
             if attr.name == "to" and attr.i == onnx.TensorProto.FLOAT:
                 attr.i = onnx.TensorProto.FLOAT16
                 break

modelopt/onnx/autocast/precisionconverter.py (1)

921-923: Prefer public onnx_utils APIs over underscored helpers across modules.

Using _get_tensor_type_by_name, _bypass_cast_node, and _is_same_type_cast from another module couples this class to private implementation details. Exposing public wrappers in modelopt/onnx/utils.py would make this boundary safer.

Also applies to: 942-942, 1097-1102

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/onnx/autocast/precisionconverter.py` around lines 921 - 923, The
code currently calls private helpers onnx_utils._get_tensor_type_by_name,
onnx_utils._bypass_cast_node, and onnx_utils._is_same_type_cast from
precisionconverter.py; replace these calls with public wrapper APIs (e.g.,
get_tensor_type_by_name, bypass_cast_node, is_same_type_cast) exported from
modelopt/onnx/utils.py and update precisionconverter.py to call those public
names (also update the other locations that use the underscored helpers around
lines referenced). Add the small public wrapper implementations in
modelopt/onnx/utils.py that delegate to the existing private functions so other
modules use the stable public API and then run tests to ensure behavior is
unchanged.

modelopt/torch/_deploy/utils/torch_onnx.py (2)

576-576: ⚠️ Potential issue | 🟠 Major

Rebuild model_metadata after the last graph rewrite.

Starting with quantize_weights(), this function mutates node names, tensor dtypes, and Q/DQ structure after model_metadata has already been captured above. The returned metadata can therefore describe a different graph than the bytes written to disk.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/_deploy/utils/torch_onnx.py` at line 576, model_metadata is
captured before quantize_weights mutates the ONNX graph (node names, tensor
dtypes, Q/DQ structure), so the saved metadata can be out of sync with
onnx_opt_graph; after the final graph rewrite (the call to quantize_weights that
returns onnx_opt_graph) re-run the metadata extraction routine (the same
function that produced model_metadata earlier) to rebuild model_metadata from
the mutated onnx_opt_graph so the metadata matches the bytes written to disk;
update any downstream uses to reference the new model_metadata variable produced
after quantize_weights.

---

581-592: ⚠️ Potential issue | 🟠 Major

The FP16 path is still “autocast + convert_float_to_float16()”.

With this new rewrite block, weights_dtype="fp16" still traces under torch.autocast("cuda") earlier in the function, so FP16 export is not actually using convert_float_to_float16() instead of autocast. That makes the exported graph depend on both mechanisms and undermines the stated pipeline change.

🛠️ Suggested earlier-function change

-    use_torch_autocast = not (
-        is_fp4_quantized(model) or is_mxfp8_quantized(model) or weights_dtype == "fp32"
-    )
+    use_torch_autocast = weights_dtype == "bf16" and not (
+        is_fp4_quantized(model) or is_mxfp8_quantized(model)
+    )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/_deploy/utils/torch_onnx.py` around lines 581 - 592, The
current FP16 branch still runs under torch.autocast earlier, so the export uses
both autocast and convert_float_to_float16; modify the export flow so when
weights_dtype == "fp16" you do NOT run the model export under
torch.autocast("cuda") (or skip the autocast context) so the graph is produced
in full-precision then transformed only by convert_float_to_float16 and
change_casts_to_fp16; locate the autocast usage earlier in this module
(torch.autocast or the export context manager) and add a conditional to bypass
it when weights_dtype == "fp16", ensuring
convert_float_to_float16/remove_redundant_casts are the sole FP16
transformations.

modelopt/onnx/export/fp8_exporter.py (1)

121-147: ⚠️ Potential issue | 🟠 Major

Reject native FP8 Q/DQ export below opset 19.

This conversion emits native QuantizeLinear/DequantizeLinear with FLOAT8E4M3FN and saturate, but it still doesn't guard against onnx_opset < 19. Exporting FP8 with a lower opset will silently produce an invalid model instead of failing early or upgrading the opset.

🛠️ Suggested guard

     def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+        opset = next(
+            (op.version for op in onnx_model.opset_import if op.domain in ("", "ai.onnx")),
+            0,
+        )
+        if opset < 19:
+            raise ValueError("Native FP8 ONNX Q/DQ requires ai.onnx opset >= 19.")
+
         logger.info("Post-processing FP8 quantized model")
         graph = gs.import_onnx(onnx_model)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/onnx/export/fp8_exporter.py` around lines 121 - 147, The conversion
loop for TRT_FP8QuantizeLinear/TRT_FP8DequantizeLinear must guard against onnx
opset < 19; detect the model/export opset (e.g., an existing variable or by
examining graph.opset or a passed opset parameter) before performing the
conversions in the loops that change node.op to
"QuantizeLinear"/"DequantizeLinear" and set FLOAT8E4M3FN/saturate, and if the
opset is less than 19 either raise a clear exception or upgrade the opset to >=
19 before modifying nodes (apply this check where you manipulate node.op and
node.attrs in the FP8 exporter function that iterates graph.nodes).

modelopt/onnx/utils.py (3)

1420-1435: ⚠️ Potential issue | 🟠 Major

Missing continue after _is_sequential_cast can cause double processing.

A node matching both _is_sequential_cast and _is_foldable_constant_cast_pattern (e.g., Constant → Cast → Cast) will be processed by both blocks, leading to duplicate entries in nodes_to_remove and potential errors during removal.

🛠️ Proposed fix

             # Find sequential casts that don't change precision
             if _is_sequential_cast(onnx_model, node):
                 nodes_to_remove.append(node)
                 _bypass_cast_node(onnx_model, node)
                 logger.debug(f"Found removable double-cast: {node.name}")
+                continue

             # Find foldable Constant -> Cast. Initializers are handled by _convert_initializers.
             if _is_foldable_constant_cast_pattern(onnx_model, node):

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/onnx/utils.py` around lines 1420 - 1435, The sequential-cast branch
(_is_sequential_cast) can fall through and be reprocessed by the foldable
Constant->Cast branch, causing duplicate entries in nodes_to_remove and
double-modification of the same node; after handling a sequential cast (where
you call _bypass_cast_node and log), add a control-flow break (e.g., a continue)
to skip further checks for that node so it isn't also processed by
_is_foldable_constant_cast_pattern, ensuring nodes_to_remove, _bypass_cast_node,
_convert_constant_values, get_producer_nodes and the node aren't acted on twice.

---

1516-1528: ⚠️ Potential issue | 🟠 Major

Using any() can incorrectly change shared Cast outputs affecting non-target consumers.

When a Cast node's output feeds multiple consumers and only some are in target_op_types, changing the Cast to FP16 affects all consumers, potentially breaking non-target branches.

🛠️ Proposed fix - only change if ALL consumers are target ops

         # Check if this Cast outputs to a target op type
         cast_output = node.output[0]
         consumers = tensor_to_consumers.get(cast_output, [])
-        feeds_target = any(c.op_type in target_op_types for c in consumers)
+        feeds_target = bool(consumers) and all(c.op_type in target_op_types for c in consumers)

         if not feeds_target:
             continue

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/onnx/utils.py` around lines 1516 - 1528, The current logic uses
any() and flips a Cast to FP16 even when some consumers are non-target, which
can break those branches; change the feeds_target condition to only be true when
there is at least one consumer and ALL consumers are in target_op_types (i.e.,
replace feeds_target = any(...) with feeds_target = bool(consumers) and
all(c.op_type in target_op_types for c in consumers)), then proceed with the
existing node.attribute loop that checks attr.name == "to" and attr.i to change
FLOAT to FLOAT16.

---

1262-1277: ⚠️ Potential issue | 🟡 Minor

Function does not handle Constant node outputs that lack value_info entries.

When a Constant node produces an output tensor that isn't registered in value_info, initializer, input, or output, this function raises an exception. This can occur before _is_foldable_constant_cast_pattern gets a chance to handle it.

Consider checking node outputs for Constant nodes:

🛠️ Proposed fix

 def _get_tensor_type_by_name(model: onnx.ModelProto, tensor_name: str):
     """Get the tensor element type. Searches value_info, initializers, inputs, and outputs."""
     for vi in model.graph.value_info:
         if vi.name == tensor_name:
             return vi.type.tensor_type.elem_type
     for init in model.graph.initializer:
         if init.name == tensor_name:
             return init.data_type
     for inp in model.graph.input:
         if inp.name == tensor_name:
             return inp.type.tensor_type.elem_type
     for out in model.graph.output:
         if out.name == tensor_name:
             return out.type.tensor_type.elem_type
+    # Check Constant node outputs
+    for node in model.graph.node:
+        if node.op_type == "Constant" and tensor_name in node.output:
+            for attr in node.attribute:
+                if attr.name == "value" and attr.type == onnx.AttributeProto.TENSOR:
+                    return attr.t.data_type
+            break
     raise Exception(f"did not find tensor {tensor_name}")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/onnx/utils.py` around lines 1262 - 1277, _get_tensor_type_by_name
currently misses Constant node outputs that aren't listed in
value_info/initializer/input/output; update it to also scan model.graph.node for
nodes with op_type == "Constant" whose output list contains tensor_name, then
find the node AttributeProto with name "value" (the TensorProto stored on the
Constant) and return its data_type (attr.t.data_type). Keep the existing checks
first (value_info/initializer/input/output), add the Constant-node check before
raising the Exception, and otherwise keep the same error behavior.

modelopt/onnx/utils.py (1)

1289-1294: Redundant None check on list comprehension result.

input_types is always a list (from the list comprehension), so input_types is not None is always True. This check can be simplified.

♻️ Proposed fix

 def _is_same_type_cast(model: onnx.ModelProto, node: onnx.NodeProto) -> bool:
     assert node.op_type == "Cast"
     input_types = [_get_tensor_type_by_name(model, inp) for inp in node.input]
     output_type = get_cast_to_type(node)
-    return all(inp_type == output_type for inp_type in input_types) and input_types is not None
+    return bool(input_types) and all(inp_type == output_type for inp_type in input_types)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/onnx/utils.py` around lines 1289 - 1294, The function
_is_same_type_cast has a redundant "input_types is not None" check because
input_types is always a list from the list comprehension; update the return
statement in _is_same_type_cast (referencing variables input_types and
output_type) to simply return all(inp_type == output_type for inp_type in
input_types), removing the unnecessary None check.

modelopt/torch/_deploy/utils/torch_onnx.py (1)

573-576: model_metadata created before graph mutations may have stale onnx_node_names.

Metadata is captured at line 573-575 before quantize_weights, FP16 conversion, and remove_redundant_casts. While I/O names and shapes should remain stable, onnx_node_names will be stale after remove_redundant_casts removes Cast nodes.

If onnx_node_names is used for validation or debugging, consider moving metadata creation after line 606 (after all mutations).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/_deploy/utils/torch_onnx.py` around lines 573 - 576, The code
creates model_metadata via create_model_metadata(...) before performing graph
mutations, which can leave onnx_node_names stale; move the model_metadata =
create_model_metadata(tree_spec_input, tree_spec_output, input_none_names,
onnx_opt_graph, model) call to after the transformations (after
quantize_weights, FP16 conversion/convert_model_to_fp16, and
remove_redundant_casts) so that onnx_opt_graph reflects the final node set;
ensure you reference the same onnx_opt_graph and model when recreating metadata
so onnx_node_names are accurate for validation/debugging.